package com.felk.kafka.java;

import com.alibaba.fastjson2.JSON;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.NoArgsConstructor;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.serialization.StringSerializer;

import java.util.Map;
import java.util.Properties;
import java.util.concurrent.TimeUnit;

// 实现自定义分区，将 MAN 的数据分到一个区，WOMAN的分到另一个区，都不是根据hash分区
public class E4_KafkaProducerPartition {
    private final static String topic = "kafkabase_test";

    public static void main(String[] args) {
        try {
            String messageStr = "这是第%s条数据";
            new InnerProducer(topic).repeat_send(messageStr, 1000);
            TimeUnit.SECONDS.sleep(5);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static class InnerProducer {
        private KafkaProducer kafkaProducer;

        private String topic;

        public InnerProducer(String topic) {
            Properties props = new Properties();
            props.put("bootstrap.servers", "192.168.56.101:9092");
            props.put("acks", "all");
            props.put("retries", 0);
            props.put("batch.size", 16384);
            props.put("key.serializer", StringSerializer.class.getName());
            props.put("value.serializer", FastJsonSerializable.class.getName());
//            props.put("value.serializer", StringSerializer.class.getName());
            // 设置分区器
            props.put("partitioner.class", MyPartitioner.class.getName());
//            props.put("value.serializer", StringSerializer.class.getName());
            this.kafkaProducer = new KafkaProducer<String, String>(props);
            this.topic = topic;
        }

        public void repeat_send(String msg, int num) {
            for (int i = 0; i < num; i++) {
                String _msg = String.format(msg, i);

                Record.TypeEnum type = Record.TypeEnum.MAN;
                if (i % 2 == 0) {
                    type = Record.TypeEnum.WOMAN;
                }
                String userId = "userid_" + (i % 4);
                Record record = Record.of(type, userId, _msg);
                send(record);
                System.out.println(record);
            }
        }

        public void send(Record record) {
            kafkaProducer.send(new ProducerRecord(topic, record.getUserId(), record));
        }
    }


    public static class MyPartitioner implements Partitioner {

        @Override
        public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
            Integer partitionCount = cluster.partitionCountForTopic(topic);
            if (partitionCount < 2) {
                throw new IllegalArgumentException(topic + "对应的分区数" + partitionCount + "小于2，不满足系统要求，请配置partition数大于2");
            }
            if (value instanceof Record) {
                Record record = (Record) value;
                if (record.getType().equals(Record.TypeEnum.MAN)) {
                    return 0;
                } else if (record.getType().equals(Record.TypeEnum.WOMAN)) {
                    return 1;
                }
                return key.toString().hashCode() % (partitionCount - 1);
            }
            return 0;
        }

        @Override
        public void close() {

        }

        @Override
        public void configure(Map<String, ?> configs) {

        }
    }

    public static class FastJsonSerializable implements Serializer<Record> {

        @Override
        public byte[] serialize(String topic, Record data) {
            return JSON.toJSONBytes(data);
        }
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    public static class Record {
        private TypeEnum type;
        private String userId;
        private String msg;

        public static Record of(TypeEnum type, String userId, String msg) {
            return new Record(type, userId, msg);
        }

        @Getter
        @AllArgsConstructor
        @NoArgsConstructor
        public static enum TypeEnum {
            MAN(1), WOMAN(2);
            private int code;
        }
    }

}